"""
IO工具模块，实现文件读写和数据转换功能。
"""
import os
import json
import csv
import glob
import pandas as pd
from typing import Dict, List, Optional, Any, Union, Tuple
import torch

def save_json(data: Any, filepath: str, indent: int = 4) -> None:
    """
    保存数据为JSON文件
    
    Args:
        data: 要保存的数据
        filepath: 文件路径
        indent: 缩进空格数
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    with open(filepath, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=indent)

def load_json(filepath: str) -> Any:
    """
    从JSON文件加载数据
    
    Args:
        filepath: 文件路径
    
    Returns:
        加载的数据
    """
    with open(filepath, 'r', encoding='utf-8') as f:
        return json.load(f)

def save_csv(data: List[Dict[str, Any]], filepath: str) -> None:
    """
    保存数据为CSV文件
    
    Args:
        data: 要保存的数据列表
        filepath: 文件路径
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    if not data:
        return
    
    # 获取所有字段名
    fieldnames = list(data[0].keys())
    
    with open(filepath, 'w', newline='', encoding='utf-8') as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(data)

def load_csv(filepath: str) -> List[Dict[str, Any]]:
    """
    从CSV文件加载数据
    
    Args:
        filepath: 文件路径
    
    Returns:
        加载的数据列表
    """
    with open(filepath, 'r', newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        return list(reader)

def csv_to_tensorboard(csv_dir: str, output_dir: str, tag_prefix: str = "") -> None:
    """
    将CSV数据转换为TensorBoard格式
    
    Args:
        csv_dir: CSV文件目录
        output_dir: 输出目录
        tag_prefix: 标签前缀
    """
    from torch.utils.tensorboard import SummaryWriter
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 创建SummaryWriter
    writer = SummaryWriter(output_dir)
    
    # 查找所有CSV文件
    csv_files = glob.glob(os.path.join(csv_dir, "*.csv"))
    
    for csv_file in csv_files:
        # 获取文件名（不包含扩展名），作为标签名
        tag_name = os.path.splitext(os.path.basename(csv_file))[0]
        if tag_prefix:
            tag_name = f"{tag_prefix}/{tag_name}"
        
        # 读取CSV文件
        try:
            df = pd.read_csv(csv_file)
            
            # 检查是否存在'step'列
            if 'step' not in df.columns:
                continue
            
            # 遍历所有数值列，添加到TensorBoard
            for column in df.columns:
                if column == 'step':
                    continue
                
                # 检查是否为数值类型
                if pd.api.types.is_numeric_dtype(df[column]):
                    for _, row in df.iterrows():
                        writer.add_scalar(f"{tag_name}/{column}", row[column], row['step'])
        except:
            # 忽略读取错误
            pass
    
    writer.close()

def save_tensor(tensor: torch.Tensor, filepath: str) -> None:
    """
    保存张量到文件
    
    Args:
        tensor: 要保存的张量
        filepath: 文件路径
    """
    # 确保目录存在
    os.makedirs(os.path.dirname(os.path.abspath(filepath)), exist_ok=True)
    
    torch.save(tensor, filepath)

def load_tensor(filepath: str, device: str = "cpu") -> torch.Tensor:
    """
    从文件加载张量
    
    Args:
        filepath: 文件路径
        device: 设备
    
    Returns:
        加载的张量
    """
    return torch.load(filepath, map_location=device)

def ensure_directory(directory: str) -> str:
    """
    确保目录存在，如果不存在则创建
    
    Args:
        directory: 目录路径
    
    Returns:
        目录路径
    """
    os.makedirs(directory, exist_ok=True)
    return directory

def get_latest_file(directory: str, pattern: str = "*") -> Optional[str]:
    """
    获取目录中最新的文件
    
    Args:
        directory: 目录路径
        pattern: 匹配模式
    
    Returns:
        最新文件的路径，如果没有文件则返回None
    """
    files = glob.glob(os.path.join(directory, pattern))
    if not files:
        return None
    
    return max(files, key=os.path.getctime) 